package hamaster.jpbc.crypt.aes;

import static hamaster.jpbc.crypt.aes.Rijndael.AES_BLOCK_SIZE;
import static hamaster.jpbc.crypt.aes.Rijndael.DIR_BOTH;
import hamaster.jpbc.crypt.Crypt;

/**
 * AES实现<br />
 * 支持ECB, CBC, OFB, CFB四种操作模式和128, 192, 256位密钥<br />
 * 只采用0填充模式<br />
 * 这个实现<b style="color:red">不是</b>线程安全的
 * @author <a href="mailto:wangyeee@gmail.com">Wang Ye</a>
 * @see OperationMode
 */
public class AES implements Crypt {

	private Rijndael rijndael;
	private int blockSize;
	private OperationMode operationMode;
	private volatile byte[] initializationVector;

	public AES() {
		rijndael = new Rijndael();
		operationMode = OperationMode.CBC;
		initializationVector = new byte[AES_BLOCK_SIZE];
	}

	@Override
	public byte[] encrypt(byte[] plain) {
		if (plain.length != blockSize)
			throw new IllegalArgumentException("Invalid block size:" + plain.length);
		switch (operationMode) {
		case ECB:
			return encryptECB(plain);
		case CBC:
			return encryptCBC(plain);
		case CFB:
			return decryptCBC(plain);
		case OFB:
			return cryptOFB(plain);
		default:
			return null;
		}
	}

	@Override
	public byte[] decrypt(byte[] cipher) {
		if (cipher.length != this.blockSize)
			throw new IllegalArgumentException("Invalid block size:" + cipher.length);
		switch (operationMode) {
		case ECB:
			return decryptECB(cipher);
		case CBC:
			return decryptCBC(cipher);
		case CFB:
			return encryptCBC(cipher);
		case OFB:
			return cryptOFB(cipher);
		default:
			throw new IllegalArgumentException("Invalid OperationMode:" + operationMode);
		}
	}

	@Override
	public void setKey(byte[] key) {
		rijndael.makeKey(key, key.length * 8, DIR_BOTH);
	}

	@Override
	public void setBlockSize(int blockSize) {
		if (blockSize % AES_BLOCK_SIZE != 0)
			throw new IllegalArgumentException("Illegal block size:" + blockSize);
		this.blockSize = blockSize;
	}

	public void setOperationMode(OperationMode operationMode) {
		this.operationMode = operationMode;
	}

	public OperationMode getOperationMode() {
		return operationMode;
	}

	public void setInitializationVector(byte[] initializationVector) {
		if (initializationVector.length != AES_BLOCK_SIZE)
			throw new IllegalArgumentException("invalid initialization vector length:" + initializationVector.length);
		System.arraycopy(initializationVector, 0, this.initializationVector, 0, AES_BLOCK_SIZE);
	}

	public byte[] getInitializationVector() {
		return initializationVector;
	}

	/**
	 * data ^= iv
	 * @param iv
	 * @param data
	 */
	final private static void xorArray(byte[] iv, byte[] data) {
		if (iv == null || data == null)
			return;
		if (iv.length != data.length)
			return;
		for (int i = 0; i < data.length; i++) {
			data[i] ^= iv[i];
		}
	}

	/**
	 * CBC模式
	 */
	private byte[] encryptCBC(byte[] pt) {
		byte[] plaint = new byte[AES_BLOCK_SIZE];
		byte[] result = new byte[blockSize];
		int bytesCrypted = 0;
		while (bytesCrypted < blockSize) {
			System.arraycopy(pt, bytesCrypted, plaint, 0, AES_BLOCK_SIZE);
			xorArray(initializationVector, plaint);
			byte[] ciphert = rijndael.encrypt(plaint);
			System.arraycopy(ciphert, 0, result, bytesCrypted, AES_BLOCK_SIZE);
			System.arraycopy(ciphert, 0, initializationVector, 0, AES_BLOCK_SIZE);
			bytesCrypted += AES_BLOCK_SIZE;
		}
		return result;
	}

	/**
	 * CBC模式
	 */
	private byte[] decryptCBC(byte[] ct) {
		byte[] ciphert = new byte[AES_BLOCK_SIZE];
		byte[] result = new byte[blockSize];
		int bytesCrypted = 0;
		while (bytesCrypted < blockSize) {
			System.arraycopy(ct, bytesCrypted, ciphert, 0, AES_BLOCK_SIZE);
			byte[] plaint = rijndael.decrypt(ciphert);
			xorArray(initializationVector, plaint);
			System.arraycopy(plaint, 0, result, bytesCrypted, AES_BLOCK_SIZE);
			System.arraycopy(ciphert, 0, initializationVector, 0, AES_BLOCK_SIZE);
			bytesCrypted += AES_BLOCK_SIZE;
		}
		return result;
	}

	/**
	 * AES加密 ECB模式
	 */
	private byte[] encryptECB(byte[] pt) {
		byte[] plaint = new byte[AES_BLOCK_SIZE];
		byte[] result = new byte[blockSize];
		int bytesCrypted = 0;
		while (bytesCrypted < blockSize) {
			System.arraycopy(pt, bytesCrypted, plaint, 0, AES_BLOCK_SIZE);
			byte[] ciphert = rijndael.encrypt(plaint);
			System.arraycopy(ciphert, 0, result, bytesCrypted, AES_BLOCK_SIZE);
			bytesCrypted += AES_BLOCK_SIZE;
		}
		return result;
	}

	/**
	 * AES解密 ECB模式
	 */
	private byte[] decryptECB(byte[] ct) {
		byte[] ciphert = new byte[AES_BLOCK_SIZE];
		byte[] result = new byte[blockSize];
		int bytesCrypted = 0;
		while (bytesCrypted < blockSize) {
			System.arraycopy(ct, bytesCrypted, ciphert, 0, AES_BLOCK_SIZE);
			byte[] plaint = rijndael.decrypt(ciphert);
			System.arraycopy(plaint, 0, result, bytesCrypted, AES_BLOCK_SIZE);
			bytesCrypted += AES_BLOCK_SIZE;
		}
		return result;
	}

	/**
	 * OFB 模式 加密和解密一样
	 */
	private byte[] cryptOFB(byte[] data) {
		byte[] plaint = new byte[AES_BLOCK_SIZE];
		byte[] result = new byte[blockSize];
		int bytesCrypted = 0;
		while (bytesCrypted < blockSize) {
			byte[] ciphert = rijndael.encrypt(initializationVector);
			System.arraycopy(ciphert, 0, initializationVector, 0, AES_BLOCK_SIZE);
			System.arraycopy(data, bytesCrypted, plaint, 0, AES_BLOCK_SIZE);
			xorArray(plaint, ciphert);
			System.arraycopy(ciphert, 0, result, bytesCrypted, AES_BLOCK_SIZE);
			bytesCrypted += AES_BLOCK_SIZE;
		}
		return result;
	}
}
